import sys
import pandas as pd
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout,
    QPushButton, QLabel, QFileDialog, QComboBox,
    QLineEdit, QMessageBox, QDialog, QListWidget,
    QDialogButtonBox, QAbstractItemView, QFrame
)
from PyQt5.QtCore import Qt, QThread, QObject, pyqtSignal
import pandas.api.types as ptypes


# MergeWorker 类 (核心算法已完全重写)
class MergeWorker(QObject):
    finished = pyqtSignal()
    error = pyqtSignal(str)
    progress = pyqtSignal(str)

    def __init__(self, df_a, list_of_dfs_b, key_a, key_b, value_col_b, new_col_name, save_filepath):
        super().__init__()
        self.df_a = df_a
        self.list_of_dfs_b = list_of_dfs_b
        self.key_a = key_a
        self.key_b = key_b
        self.value_col_b = value_col_b
        self.new_col_name = new_col_name
        self.save_filepath = save_filepath

    def run(self):
        try:
            self.progress.emit("正在初始化...")
            self.df_a[self.new_col_name] = None
            self.df_a[self.key_a] = self.df_a[self.key_a].astype(str)

            for i, df_b_sheet in enumerate(self.list_of_dfs_b):
                self.progress.emit(f"正在处理第 {i + 1} 个数据源...")

                if self.key_b not in df_b_sheet.columns or self.value_col_b not in df_b_sheet.columns:
                    continue

                df_b_sheet[self.key_b] = df_b_sheet[self.key_b].astype(str)

                # --- 【核心修改点：根据列类型选择不同逻辑】 ---

                # 1. 自动检测要获取的列（value_col_b）是否为数值类型
                is_numeric_col = ptypes.is_numeric_dtype(df_b_sheet[self.value_col_b])

                if is_numeric_col:
                    # --- 情况A：如果是数值列（如金额），执行“分组求和” ---
                    self.progress.emit(f"检测到数值列，执行加总匹配...")

                    # 确保B表的值列是数字类型
                    df_b_sheet[self.value_col_b] = pd.to_numeric(df_b_sheet[self.value_col_b], errors='coerce').fillna(
                        0)

                    # 按B表的键进行分组，并对值列求和
                    mapping = df_b_sheet.groupby(self.key_b)[self.value_col_b].sum()

                else:
                    # --- 情况B：如果是文本列（如姓名），执行“优先查找” ---
                    self.progress.emit(f"检测到文本列，执行优先查找匹配...")

                    # 去除重复的键，只保留第一个出现的，然后创建映射
                    mapping = df_b_sheet.drop_duplicates(subset=[self.key_b]).set_index(self.key_b)[self.value_col_b]

                # --- 【修改结束】---

                # 后续的填充逻辑保持不变，它会自动使用上面生成的正确mapping
                self.df_a.loc[self.df_a[self.new_col_name].isnull(), self.new_col_name] = self.df_a[self.key_a].map(
                    mapping)

            self.progress.emit("正在保存文件...")
            self.df_a[self.new_col_name].fillna('', inplace=True)

            if self.save_filepath.endswith('.xlsx'):
                with pd.ExcelWriter(self.save_filepath, engine='xlsxwriter') as writer:
                    self.df_a.to_excel(writer, sheet_name='匹配结果', index=False)
                    workbook = writer.book
                    worksheet = writer.sheets['匹配结果']
                    text_format = workbook.add_format({'num_format': '@'})
                    key_a_col_idx = self.df_a.columns.get_loc(self.key_a)
                    if self.new_col_name in self.df_a.columns:
                        new_col_idx = self.df_a.columns.get_loc(self.new_col_name)
                        worksheet.set_column(new_col_idx, new_col_idx, None, text_format)
                    worksheet.set_column(key_a_col_idx, key_a_col_idx, None, text_format)
            else:
                self.df_a.to_csv(self.save_filepath, index=False, encoding='utf_8_sig')

            self.finished.emit()
        except Exception as e:
            self.error.emit(str(e))


# SheetSelectionDialog 类 (无需改动)
class SheetSelectionDialog(QDialog):
    # ... 此处代码与上一版完全相同 ...
    def __init__(self, sheet_names, mode='multiple', parent=None):
        super().__init__(parent)
        self.setWindowTitle("选择工作表")
        self.setMinimumWidth(300)
        layout = QVBoxLayout(self)
        layout.addWidget(QLabel("请选择要加载的工作表:"))
        self.list_widget = QListWidget()
        if mode == 'single':
            self.list_widget.setSelectionMode(QAbstractItemView.SingleSelection)
        else:
            self.list_widget.setSelectionMode(QAbstractItemView.ExtendedSelection)
        self.list_widget.addItems(sheet_names)
        layout.addWidget(self.list_widget)
        if mode == 'single':
            self.list_widget.itemDoubleClicked.connect(self.accept)
        button_box = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        button_box.accepted.connect(self.accept)
        button_box.rejected.connect(self.reject)
        layout.addWidget(button_box)

    def get_selected_sheets(self):
        return [item.text() for item in self.list_widget.selectedItems()]


# 主应用类 (已更新)
class CompactApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("优先顺序查找工具 V6.0")
        self.setFixedSize(299, 458)
        self.df_a = None
        self.list_of_dfs_b = None  # B表现在是一个DataFrame列表
        self.thread = None
        self.worker = None

        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        main_layout = QVBoxLayout(main_widget)
        main_layout.setSpacing(6)
        main_layout.setContentsMargins(10, 10, 10, 10)

        # UI 布局 (移除了去重选框)
        main_layout.addWidget(QLabel("<b>步骤 1: 加载主表 (A)</b>"))
        self.btn_load_a = QPushButton("加载表A (单Sheet)")
        self.label_a_info = QLabel("未加载")
        self.label_a_info.setStyleSheet("color: gray;")
        self.combo_key_a = QComboBox()
        main_layout.addWidget(self.btn_load_a)
        main_layout.addWidget(self.label_a_info)
        main_layout.addWidget(QLabel("表A对比键:"))
        main_layout.addWidget(self.combo_key_a)

        main_layout.addStretch(1)

        main_layout.addWidget(QLabel("<b>步骤 2: 加载数据源 (B)</b>"))
        self.btn_load_b = QPushButton("加载表B (可多Sheet, 有顺序)")
        self.label_b_info = QLabel("未加载")
        self.label_b_info.setStyleSheet("color: gray;")
        self.combo_key_b = QComboBox()
        self.combo_value_b = QComboBox()
        main_layout.addWidget(self.btn_load_b)
        main_layout.addWidget(self.label_b_info)
        main_layout.addWidget(QLabel("表B对比键:"))
        main_layout.addWidget(self.combo_key_b)
        main_layout.addWidget(QLabel("从表B获取此列:"))
        main_layout.addWidget(self.combo_value_b)

        main_layout.addStretch(1)

        main_layout.addWidget(QLabel("<b>步骤 3: 命名新列</b>"))
        self.new_col_name_input = QLineEdit()
        self.new_col_name_input.setPlaceholderText("例如: 匹配到的下单人")
        main_layout.addWidget(self.new_col_name_input)

        main_layout.addStretch(10)
        separator = QFrame()
        separator.setFrameShape(QFrame.HLine)
        separator.setFrameShadow(QFrame.Sunken)
        main_layout.addWidget(separator)
        self.btn_run_save = QPushButton("执行并保存")
        self.btn_run_save.setStyleSheet("background-color: #4CAF50; color: white; font-weight: bold; padding: 5px;")
        main_layout.addWidget(self.btn_run_save)

        self.btn_load_a.clicked.connect(lambda: self.load_table('a'))
        self.btn_load_b.clicked.connect(lambda: self.load_table('b'))
        self.btn_run_save.clicked.connect(self.run_and_save)

    def load_table(self, table_id):
        # load_table 函数 (B表加载逻辑已修改)
        filepath, _ = QFileDialog.getOpenFileName(
            self, "选择文件", "", "Excel 文件 (*.xlsx *.xls);;CSV 文件 (*.csv)")
        if not filepath: return
        try:
            if table_id == 'a':
                # ... A表加载逻辑无改动 ...
                if not filepath.endswith(('.xlsx', '.xls')):
                    QMessageBox.warning(self, "格式错误", "表A必须是Excel文件才能选择Sheet。")
                    return
                xls = pd.ExcelFile(filepath)
                dialog = SheetSelectionDialog(xls.sheet_names, mode='single', parent=self)
                if dialog.exec_() == QDialog.Accepted:
                    selected_sheets = dialog.get_selected_sheets()
                    if not selected_sheets: return
                    sheet_name = selected_sheets[0]
                    self.df_a = pd.read_excel(filepath, sheet_name=sheet_name)
                    self.combo_key_a.clear()
                    self.combo_key_a.addItems(self.df_a.columns)
                    self.label_a_info.setText(f"{filepath.split('/')[-1]} [{sheet_name}]")
                    self.label_a_info.setStyleSheet("color: green;")
            else:  # B表加载逻辑修改
                if not filepath.endswith(('.xlsx', '.xls')):
                    QMessageBox.warning(self, "格式错误", "表B必须是Excel文件才能选择多个Sheet。")
                    return
                xls = pd.ExcelFile(filepath)
                dialog = SheetSelectionDialog(xls.sheet_names, mode='multiple', parent=self)
                if dialog.exec_() == QDialog.Accepted:
                    selected_sheets = dialog.get_selected_sheets()
                    if not selected_sheets: return

                    # 将B表加载为DataFrame的列表，保持顺序
                    self.list_of_dfs_b = [pd.read_excel(filepath, sheet_name=s) for s in selected_sheets]

                    # 下拉框用第一个Sheet的内容填充
                    first_sheet_df = self.list_of_dfs_b[0]
                    self.combo_key_b.clear()
                    self.combo_value_b.clear()
                    self.combo_key_b.addItems(first_sheet_df.columns)
                    self.combo_value_b.addItems(first_sheet_df.columns)
                    self.label_b_info.setText(f"{filepath.split('/')[-1]} (已加载 {len(selected_sheets)}个Sheet)")
                    self.label_b_info.setStyleSheet("color: green;")
        except Exception as e:
            QMessageBox.critical(self, "加载错误", f"无法加载文件:\n{e}")

    def run_and_save(self):
        # run_and_save 函数 (已更新)
        if self.df_a is None or self.list_of_dfs_b is None:
            QMessageBox.warning(self, "缺少数据", "请先加载表A和表B。")
            return
        key_a = self.combo_key_a.currentText()
        key_b = self.combo_key_b.currentText()
        value_col_b = self.combo_value_b.currentText()
        new_col_name = self.new_col_name_input.text().strip()
        if not all([key_a, key_b, value_col_b, new_col_name]):
            QMessageBox.warning(self, "输入不完整", "请确保所有下拉框都已选择，并为新列命名。")
            return
        save_filepath, _ = QFileDialog.getSaveFileName(
            self, "选择保存位置", "", "Excel 文件 (*.xlsx);;CSV 文件 (*.csv)")
        if not save_filepath: return

        self.set_ui_enabled(False)
        self.thread = QThread()
        self.worker = MergeWorker(
            df_a=self.df_a.copy(),
            list_of_dfs_b=[df.copy() for df in self.list_of_dfs_b],  # 传递B表列表
            key_a=key_a, key_b=key_b, value_col_b=value_col_b,
            new_col_name=new_col_name, save_filepath=save_filepath
        )
        self.worker.moveToThread(self.thread)
        self.thread.started.connect(self.worker.run)
        self.worker.finished.connect(self.on_task_finished)
        self.worker.error.connect(self.on_task_error)
        self.worker.progress.connect(self.update_status)
        self.worker.finished.connect(self.thread.quit)
        self.worker.finished.connect(self.worker.deleteLater)
        self.thread.finished.connect(self.thread.deleteLater)
        self.thread.start()

    # 其他辅助函数 (set_ui_enabled, update_status,等) 无需改动
    def set_ui_enabled(self, enabled):
        self.btn_load_a.setEnabled(enabled)
        self.btn_load_b.setEnabled(enabled)
        self.btn_run_save.setEnabled(enabled)
        self.combo_key_a.setEnabled(enabled)
        self.combo_key_b.setEnabled(enabled)
        self.combo_value_b.setEnabled(enabled)
        self.new_col_name_input.setEnabled(enabled)
        if not enabled:
            self.btn_run_save.setText("正在处理...")
        else:
            self.btn_run_save.setText("执行并保存")

    def update_status(self, message):
        self.label_b_info.setText(message)
        self.label_b_info.setStyleSheet("color: blue;")

    def on_task_finished(self):
        self.set_ui_enabled(True)
        QMessageBox.information(self, "成功", "处理完成！\n新文件已保存。")
        self.label_b_info.setText("任务已完成")
        self.label_b_info.setStyleSheet("color: green;")

    def on_task_error(self, error_message):
        self.set_ui_enabled(True)
        QMessageBox.critical(self, "处理失败", f"发生错误: {error_message}")
        self.label_b_info.setText("任务失败")
        self.label_b_info.setStyleSheet("color: red;")
